"""
吸收模块：药物吸收过程建模

该模块实现了药物在胃肠道中的吸收过程，包括：
- 口服给药的溶出过程
- 肠道吸收机制
- 首过效应
- 吸收动力学参数计算
- ADAM模型集成（基于Excel参数表）

作者：BlackCat@CPPO
版本：2.0.0
"""

import numpy as np
from scipy.integrate import solve_ivp
from .adam_model import ADAMModel

class AbsorptionModel:
    """
    药物吸收模型类
    
    该类模拟药物在胃肠道的吸收过程，包括溶出、吸收和转运。
    """
    
    def __init__(self, parameters):
        """
        初始化吸收模型
        
        参数:
            parameters (dict): 模型参数，包含：
                - ka: 吸收速率常数 (1/h)
                - kd: 溶出速率常数 (1/h)
                - fa: 吸收分数
                - fg: 肠道代谢分数
                - fh: 肝脏代谢分数
                - ts: 胃排空时间 (h)
                - tlag: 滞后时间 (h)
                - use_adam: 是否使用ADAM模型 (bool)
                - excel_file: Excel参数文件路径 (str)
        """
        self.params = parameters
        self.use_adam = parameters.get('use_adam', False)
        self.adam_model = None
        self.adam_params = {}
        
        # 初始化ADAM模型
        if self.use_adam:
            self.adam_model = ADAMModel(parameters)
            # 使用默认样本参数（直接从parameters获取）
            self.adam_params = parameters
        
        self.validate_parameters()
    
    def validate_parameters(self):
        """验证模型参数的有效性"""
        # 如果使用ADAM模型，不需要传统参数
        if self.use_adam:
            required_adam = ['dose', 'solubility', 'permeability']
            for param in required_adam:
                if param not in self.params:
                    self.params[param] = self.get_default_parameter(param)
        else:
            # 传统模型参数验证
            required_params = ['ka', 'kd', 'fa', 'fg', 'fh', 'ts', 'tlag']
            for param in required_params:
                if param not in self.params:
                    self.params[param] = self.get_default_parameter(param)
    
    def get_default_parameter(self, param_name: str) -> float:
        """获取默认参数值"""
        defaults = {
            'dose': 100.0,  # mg
            'solubility': 1.0,  # mg/mL
            'permeability': 1e-4,  # cm/s
            'molecular_weight': 300.0,  # g/mol
            'pKa': 7.0,
            'logP': 2.0,
            'particle_size': 50.0,  # μm
            'density': 1.2,  # g/cm³
            'ka': 1.0,  # 1/h
            'kd': 1.0,  # 1/h
            'fa': 0.9,
            'fg': 0.95,
            'fh': 0.85,
            'ts': 1.0,  # h
            'tlag': 0.1  # h
        }
        return defaults.get(param_name, 1.0)

    def calculate_absorption(self, individual_params, time_points):
        """
        计算药物吸收过程
        
        参数:
            individual_params (dict): 个体特定参数
            time_points (np.array): 时间点数组
            
        返回:
            dict: 吸收过程结果，包含：
                - rate: 吸收速率随时间变化
                - cumulative: 累积吸收量
                - fraction_absorbed: 吸收分数
                - bioavailability: 生物利用度
                - adam_data: ADAM模型结果（如果使用ADAM）
        """
        # 合并个体参数
        params = {**self.params, **individual_params}
        
        # 如果使用ADAM模型，使用ADAM计算
        if self.use_adam and self.adam_model is not None:
            return self._calculate_adam_absorption(params, time_points)
        else:
            # 使用传统吸收模型
            return self._calculate_traditional_absorption(params, time_points)
    
    def _calculate_adam_absorption(self, params, time_points):
        """使用ADAM模型计算吸收"""
        # 更新ADAM模型参数
        adam_params = {**self.params, **params}
        self.adam_model.params.update(adam_params)
        
        # 计算ADAM吸收
        adam_result = self.adam_model.calculate_adam_absorption(time_points)
        
        # 整合结果
        # 计算BCS分类
        bcs_result = self.calculate_bcs_classification(
            self.adam_model.params['solubility'],
            self.adam_model.params['permeability'],
            self.adam_model.params['dose']
        )
        
        return {
            'gastric_amount': np.zeros_like(time_points),  # ADAM模型不需要胃排空概念
            'intestinal_amount': np.sum([
                adam_result['regional_amounts'][region] 
                for region in adam_result['regional_amounts']
            ], axis=0),
            'absorbed_amount': adam_result['cumulative_absorption'],
            'absorption_rate': adam_result['total_absorption_rate'],
            'cumulative_absorption': adam_result['cumulative_absorption'],
            'rate': adam_result['total_absorption_rate'],
            'cumulative': adam_result['cumulative_absorption'],
            'fraction_absorbed': adam_result['fraction_absorbed'],
            'bioavailability': adam_result['bioavailability'],
            'plasma_concentration': adam_result['cumulative_absorption'],  # 简化为累积吸收
            'T_max': np.argmax(adam_result['cumulative_absorption']) * (time_points[1] - time_points[0]) if len(adam_result['cumulative_absorption']) > 0 else 0,
            'C_max': np.max(adam_result['cumulative_absorption']) if len(adam_result['cumulative_absorption']) > 0 else 0,
            'adam_data': adam_result,  # 包含ADAM模型的完整结果
            'bcs_classification': bcs_result
        }
    
    def _calculate_traditional_absorption(self, params, time_points):
        """使用传统模型计算吸收"""
        # 初始条件
        # [胃中药物量, 肠道中药物量, 吸收药物量]
        initial_conditions = [params.get('dose', 100), 0, 0]
        
        # 求解微分方程
        solution = solve_ivp(
            self._absorption_odes,
            [time_points[0], time_points[-1]],
            initial_conditions,
            t_eval=time_points,
            args=(params,),
            method='RK45'
        )
        
        # 提取结果
        gastric_amount = solution.y[0]
        intestinal_amount = solution.y[1]
        absorbed_amount = solution.y[2]
        
        # 计算吸收速率
        absorption_rate = np.gradient(absorbed_amount, time_points)
        
        # 计算生物利用度
        bioavailability = absorbed_amount[-1] / params.get('dose', 100)
        
        return {
            'gastric_amount': gastric_amount,
            'intestinal_amount': intestinal_amount,
            'absorbed_amount': absorbed_amount,
            'absorption_rate': absorption_rate,
            'cumulative_absorption': absorbed_amount,
            'rate': absorption_rate,
            'cumulative': absorbed_amount,
            'fraction_absorbed': absorbed_amount / params.get('dose', 100),
            'bioavailability': bioavailability,
            'plasma_concentration': absorbed_amount,  # 简化为吸收量
            'T_max': np.argmax(absorbed_amount) * (time_points[1] - time_points[0]) if len(absorbed_amount) > 0 else 0,
            'C_max': np.max(absorbed_amount) if len(absorbed_amount) > 0 else 0,
            'adam_data': None,
            'bcs_classification': None
        }
    
    def _absorption_odes(self, t, y, params):
        """
        吸收过程微分方程组
        
        状态变量:
            y[0]: 胃中药物量 (Ag)
            y[1]: 肠道中药物量 (Ai)
            y[2]: 吸收进入体循环的药物量 (Aa)
        
        参数:
            params: 模型参数字典
        """
        Ag, Ai, Aa = y
        
        # 提取参数
        ka = params['ka']
        kd = params['kd']
        fa = params['fa']
        fg = params['fg']
        fh = params['fh']
        ts = params['ts']
        tlag = params['tlag']
        
        # 胃排空函数（考虑滞后时间）
        if t < tlag:
            kg = 0
        else:
            # 使用Weibull函数描述胃排空
            kg = (1/ts) * ((t - tlag)/ts)**2 * np.exp(-((t - tlag)/ts)**3)

        # 溶出过程计算
        if Ag <= 1e-6:
            dissolution = 0
        else:
            dissolution = kd * Ag  # 溶出速率与剩余药量成正比
        
        # 肠道吸收
        dAg_dt = -dissolution - kg * Ag  # 胃排空
        dAi_dt = kg * Ag - ka * Ai  # 肠道流入 - 吸收
        dAa_dt = ka * Ai * fa * fg * fh  # 有效吸收
        
        return [dAg_dt, dAi_dt, dAa_dt]
    
    def calculate_first_order_absorption(self, ka, t, dose=100, tlag=0):
        """
        计算一级吸收模型
        
        参数:
            ka: 吸收速率常数 (1/h)
            t: 时间点数组
            dose: 给药剂量
            tlag: 滞后时间
            
        返回:
            dict: 一级吸收结果
        """
        # 考虑滞后时间
        t_eff = np.maximum(t - tlag, 0)
        
        # 一级吸收公式
        absorption_amount = dose * (1 - np.exp(-ka * t_eff))
        absorption_rate = dose * ka * np.exp(-ka * t_eff)
        
        return {
            'amount': absorption_amount,
            'rate': absorption_rate,
            'fraction': absorption_amount / dose
        }
    
    def calculate_zero_order_absorption(self, rate, duration, t, dose=100):
        """
        计算零级吸收模型
        
        参数:
            rate: 吸收速率 (单位/小时)
            duration: 吸收持续时间 (小时)
            t: 时间点数组
            dose: 总给药剂量
            
        返回:
            dict: 零级吸收结果
        """
        absorption_amount = np.minimum(rate * t, dose)
        absorption_rate = np.where(t <= duration, rate, 0)
        
        return {
            'amount': absorption_amount,
            'rate': absorption_rate,
            'fraction': absorption_amount / dose
        }
    
    def calculate_weibull_absorption(self, t, a, b, dose=100):
        """
        计算Weibull吸收模型
        
        参数:
            t: 时间点数组
            a: 形状参数
            b: 尺度参数
            dose: 给药剂量
            
        返回:
            dict: Weibull吸收结果
        """
        # Weibull累积分布函数
        absorption_amount = dose * (1 - np.exp(-((t/b)**a)))
        
        # Weibull概率密度函数（吸收速率）
        absorption_rate = dose * (a/b) * ((t/b)**(a-1)) * np.exp(-((t/b)**a))
        
        return {
            'amount': absorption_amount,
            'rate': absorption_rate,
            'fraction': absorption_amount / dose
        }
    
    def calculate_adam_bcs_integration(self, excel_file):
        """
        计算ADAM模型与BCS分类的集成分析
        
        参数:
            excel_file (str): Excel参数文件路径
            
        返回:
            dict: BCS分类和ADAM集成结果
        """
        if not self.use_adam or self.adam_model is None:
            raise ValueError("ADAM模型未启用")
        
        return self.adam_model.calculate_bcs_adam_integration(excel_file)
    
    def calculate_regional_absorption_profile(self, excel_file, time_points):
        """
        计算基于Excel参数的区域吸收曲线
        
        参数:
            excel_file (str): Excel参数文件路径
            time_points (np.array): 时间点数组
            
        返回:
            dict: 区域吸收分析结果
        """
        if not self.use_adam or self.adam_model is None:
            raise ValueError("ADAM模型未启用")
        
        # 加载Excel参数并更新
        excel_params = self.adam_model.load_excel_parameters(excel_file)
        self.adam_model.params.update(excel_params)
        
        # 计算BCS分类
        bcs_result = self.adam_model.calculate_bcs_adam_integration(excel_file)
        
        # 计算ADAM吸收
        adam_result = self.adam_model.calculate_adam_absorption(time_points)
        
        return {
            'bcs_classification': bcs_result,
            'adam_absorption': adam_result,
            'regional_analysis': {
                'stomach': adam_result['regional_absorbed']['stomach'],
                'small_intestine': np.sum([
                    adam_result['regional_absorbed']['duodenum'],
                    adam_result['regional_absorbed']['jejunum_1'],
                    adam_result['regional_absorbed']['jejunum_2'],
                    adam_result['regional_absorbed']['ileum_1'],
                    adam_result['regional_absorbed']['ileum_2'],
                    adam_result['regional_absorbed']['ileum_3']
                ], axis=0),
                'colon': adam_result['regional_absorbed']['colon']
            }
        }

    def calculate_absorption_from_formulation(self, formulation_data, t):
        """
        根据制剂数据计算吸收
        
        参数:
            formulation_data (dict): 制剂数据，包含：
                - release_profile: 释放曲线
                - solubility: 溶解度
                - permeability: 渗透性
            t: 时间点数组
            
        返回:
            dict: 基于制剂的吸收结果
        """
        # 获取释放曲线
        release_profile = formulation_data.get('release_profile', [])
        
        if len(release_profile) == 0:
            # 如果没有释放曲线，使用一级吸收
            ka = formulation_data.get('ka', 0.5)
            return self.calculate_first_order_absorption(ka, t)
        
        # 插值获取释放数据
        from scipy.interpolate import interp1d
        
        release_times = [point[0] for point in release_profile]
        release_values = [point[1] for point in release_profile]
        
        # 创建插值函数
        release_func = interp1d(
            release_times, 
            release_values, 
            kind='linear', 
            bounds_error=False, 
            fill_value=(0, release_values[-1])
        )
        
        # 计算吸收
        release_amount = release_func(t)
        absorption_rate = np.gradient(release_amount, t)
        
        return {
            'amount': release_amount,
            'rate': absorption_rate,
            'fraction': release_amount / max(release_values)
        }
    
    def calculate_bcs_classification(self, solubility, permeability, dose):
        """
        根据BCS分类系统评估吸收
        
        参数:
            solubility: 溶解度 (mg/mL)
            permeability: 渗透性 (cm/s)
            dose: 给药剂量 (mg)
            
        返回:
            dict: BCS分类结果
        """
        # BCS分类标准
        high_solubility_threshold = 250  # mL (最高剂量在pH1-7.5可溶)
        high_permeability_threshold = 2e-4  # cm/s
        
        # 计算剂量溶解度比
        dose_solubility_ratio = dose / solubility
        
        # 判断分类
        is_high_solubility = dose_solubility_ratio < high_solubility_threshold
        is_high_permeability = permeability > high_permeability_threshold
        
        bcs_class = None
        if is_high_solubility and is_high_permeability:
            bcs_class = "I"
        elif not is_high_solubility and is_high_permeability:
            bcs_class = "II"
        elif is_high_solubility and not is_high_permeability:
            bcs_class = "III"
        else:
            bcs_class = "IV"
        
        return {
            'class': bcs_class,
            'high_solubility': is_high_solubility,
            'high_permeability': is_high_permeability,
            'dose_solubility_ratio': dose_solubility_ratio,
            'expected_absorption': self._estimate_absorption_from_bcs(bcs_class)
        }
    
    def _estimate_absorption_from_bcs(self, bcs_class):
        """根据BCS分类估计吸收"""
        absorption_map = {
            'I': {'fraction': 0.9, 'variability': '低'},
            'II': {'fraction': 0.6, 'variability': '高'},
            'III': {'fraction': 0.4, 'variability': '中'},
            'IV': {'fraction': 0.2, 'variability': '高'}
        }
        return absorption_map.get(bcs_class, {'fraction': 0.5, 'variability': '中'})
    
    def calculate_food_effect(self, fed_state_params, fasted_state_params):
        """
        计算食物效应对吸收的影响
        
        参数:
            fed_state_params: 餐后状态参数
            fasted_state_params: 空腹状态参数
            
        返回:
            dict: 食物效应分析结果
        """
        # 计算参数变化
        changes = {}
        for param in ['ka', 'kd', 'fa', 'fg', 'fh']:
            if param in fed_state_params and param in fasted_state_params:
                change_ratio = fed_state_params[param] / fasted_state_params[param]
                changes[param] = {
                    'fed': fed_state_params[param],
                    'fasted': fasted_state_params[param],
                    'ratio': change_ratio
                }
        
        # 判断食物效应类型
        if changes.get('fa', {}).get('ratio', 1) > 1.2:
            effect_type = "增加吸收"
        elif changes.get('fa', {}).get('ratio', 1) < 0.8:
            effect_type = "减少吸收"
        else:
            effect_type = "无显著影响"
        
        return {
            'effect_type': effect_type,
            'parameter_changes': changes,
            'clinical_significance': self._assess_clinical_significance(changes)
        }
    
    def _assess_clinical_significance(self, changes):
        """评估临床意义"""
        significant_changes = []
        for param, data in changes.items():
            if abs(data['ratio'] - 1) > 0.2:
                significant_changes.append(param)
        
        if len(significant_changes) > 1:
            return "可能有显著临床意义"
        elif len(significant_changes) == 1:
            return "可能有轻微临床意义"
        else:
            return "无显著临床意义"